from functools import reduce

import numpy as np
import math
import pandas as pd


def load_dataset() -> tuple:
    with open('./dataset/watermelon.txt', 'r', encoding='utf-8') as dataset_file:
        record_list = list(map(lambda x: x.rstrip('\n').split(','), dataset_file.readlines()))
    # 标签列表
    label_list = record_list[0][:len(record_list[0]) - 1]
    test_set_size = 1  # 测试集大小
    # 训练集
    data_size = len(record_list)
    train_set = record_list[1:data_size - test_set_size]
    for record in train_set:
        record[6] = float(record[6])
        record[7] = float(record[7])
    # 测试集
    test_set = record_list[-1]
    test_set[6] = float(test_set[6])
    test_set[7] = float(test_set[7])
    test_set.pop()

    return train_set, test_set, label_list


# 计算(不同类别中指定连续特征的)均值、标准差
def mean_std(feature, cla):
    train_set, test_set, label_list = load_dataset()

    lst = [record[label_list.index(feature)] for record in train_set if record[-1] == cla]  # 类别为cla中指定特征feature组成的列表
    mean = round(np.mean(lst), 3)  # 均值
    std = round(np.std(lst), 3)  # 标准差

    return mean, std


# 计算先验概率P(c)
def prior_prob():
    train_set, _, _ = load_dataset()

    good_count = 0  # 初始化好瓜数量
    bad_count = 0  # 初始化坏瓜数量
    all_count = len(train_set)

    for record in train_set:  # 统计好瓜个数
        if record[-1] == "好瓜":
            good_count += 1
    for record in train_set:  # 统计坏瓜个数
        if record[-1] == "坏瓜":
            bad_count += 1

    # 计算先验概率P(c)
    good_prob = round(good_count / all_count, 3)
    bad_prob = round(bad_count / all_count, 3)

    return good_prob, bad_prob


# 计算离散属性的条件概率P(xi|c)
def prob(index, cla):
    train_set, test_set, label_list = load_dataset()
    good_count = 0  # 初始化好瓜数量
    bad_count = 0  # 初始化坏瓜数量

    for record in train_set:  # 统计好瓜个数
        if record[-1] == "好瓜":
            good_count += 1
    for record in train_set:  # 统计坏瓜个数
        if record[-1] == "坏瓜":
            bad_count += 1

    lst = [record for record in train_set if
           (record[-1] == cla) & (record[index] == test_set[index])]  # lst为cla类中第index个属性上取值为xi的样本组成的集合
    prob = round(len(lst) / (good_count if cla == "好瓜" else bad_count), 3)  # 计算条件概率

    return prob


# 计算连续属性的条件概率p(xi|c)
def conditional_prob():
    train_set, test_set, label_list = load_dataset()

    good_den_mean, good_den_std = mean_std("密度", "好瓜")  # 好瓜密度的均值、标准差
    bad_den_mean, bad_den_std = mean_std("密度", "坏瓜")  # 坏瓜密度的均值、标准差
    good_sugar_mean, good_sugar_std = mean_std("含糖率", "好瓜")  # 好瓜含糖率的均值、标准差
    bad_sugar_mean, bad_sugar_std = mean_std("含糖率", "坏瓜")  # 坏瓜含糖率的均值、标准差
    # p(密度|好瓜)
    good_density_prob = (1 / (math.sqrt(2 * math.pi) * good_den_std)) * np.exp(
        -(((test_set[label_list.index("密度")] - good_den_mean) ** 2) / (2 * (good_den_std ** 2))))
    good_density_prob = round(good_density_prob, 3)
    # p(密度|坏瓜)
    bad_density_prob = (1 / (math.sqrt(2 * math.pi) * bad_den_std)) * np.exp(
        -(((test_set[label_list.index("密度")] - bad_den_mean) ** 2) / (2 * (bad_den_std ** 2))))
    bad_density_prob = round(bad_density_prob, 3)
    # p(含糖率|好瓜)
    good_sugar_prob = (1 / (math.sqrt(2 * math.pi) * good_sugar_std)) * np.exp(
        -(((test_set[label_list.index("含糖率")] - good_sugar_mean) ** 2) / (2 * (good_sugar_std ** 2))))
    good_sugar_prob = round(good_sugar_prob, 3)
    # p(含糖率|坏瓜)
    bad_sugar_prob = (1 / (math.sqrt(2 * math.pi) * bad_sugar_std)) * np.exp(
        -(((test_set[label_list.index("含糖率")] - bad_sugar_mean) ** 2) / (2 * (bad_sugar_std ** 2))))
    bad_sugar_prob = round(bad_sugar_prob, 3)

    return good_density_prob, bad_density_prob, good_sugar_prob, bad_sugar_prob


# 预测后验概率P(c|xi)
def post_prob():
    # 计算类先验概率
    good_prob, bad_prob = prior_prob()
    # 计算离散属性的条件概率
    good_prob_list = [prob(i, '好瓜') for i in range(6)]
    bad_prob_list = [prob(i, '坏瓜') for i in range(6)]
    # 计算连续属性的条件概率
    good_density_prob, bad_density_prob, good_sugar_prob, bad_sugar_prob = conditional_prob()

    # 计算后验概率
    # 计算是好瓜的后验概率
    is_good = reduce(lambda prev, curr: prev * curr, good_prob_list, good_prob * good_density_prob * good_sugar_prob)
    # 计算是坏瓜的后验概率
    is_bad = reduce(lambda prev, curr: prev * curr, bad_prob_list, bad_prob * bad_density_prob * bad_sugar_prob)

    return is_good, is_bad


def main():
    train_set, test_set, label_list = load_dataset()

    # 导入数据集
    load_dataset()

    test_set = [test_set]
    df = pd.DataFrame(test_set, columns=label_list, index=[1])
    print(f"待测样本 X:\n{df}")

    is_good, is_bad = post_prob()
    print("\n后验概率:")
    print(f"P(好瓜|X) = {is_good}")
    print(f"P(坏瓜|X) = {is_bad}")
    print("\n预测结果:好瓜" if (is_good > is_bad) else "预测结果:坏瓜")


if __name__ == '__main__':
    main()
